from utils import *
from PCA import PCA
import torch
import numpy as np
import pickle
import time
import os
from skimage.metrics import peak_signal_noise_ratio


raw_data = plt.imread('data/lena.jpg')


def get_psnr_pr_2(filename: str):
    with open(filename, '+rb') as f:
        recon = pickle.load(f)
    result = peak_signal_noise_ratio(raw_data, recon, data_range=256)
    return result


def draw_pr_2(filename: str, n_top: int):
    with open(filename, '+rb') as f:
        recon = pickle.load(f)
    plt.imsave(f'results/recovered_lena_top_{n_top}.jpg', recon, cmap='grey')


def main_pr2(n_components: list = [10, 50, 100, 150], bins_valid=False, img_valid=False):
    if not os.path.exists('AfterPCA'):
        os.makedirs('AfterPCA')
    if not os.path.exists('results'):
        os.makedirs('results')

    if not bins_valid:
        for n_component in n_components:
            data = torch.from_numpy(raw_data).float().reshape(512, 1536)
            print('data.shape', data.shape)
            pca = PCA(n_component=n_component, device='cpu')
            pca.fit(data)
            after_pca = pca.transform(data)
            print('after_pca.shape', after_pca.shape)
            recon = pca.reconstruct(after_pca)  
            recon = recon.reshape(512, 512, 3).cpu().numpy().astype(np.uint8)
            print('recon.shape', recon.shape)
            with open(f'AfterPCA/pr2_top_{n_component}.bin', '+wb') as f:
                    pickle.dump(recon, f)

    if not img_valid:
        for n_top in n_components:
            draw_pr_2(f'AfterPCA/pr2_top_{n_top}.bin', n_top)
    
    with open('results/log_pr2.txt', '+w') as f:
        for n_top in n_components:
            psnr = get_psnr_pr_2(f'AfterPCA/pr2_top_{n_top}.bin')
            f.write(f'avg_psnr_top_{n_top}: {psnr}\n')


if __name__ == '__main__':
    raw_data = raw_data.astype(np.uint8, copy=True)
    begin = time.time()
    with torch.no_grad():
        main_pr2(bins_valid=False, img_valid=False)  # 若重建图像已打好包，可把 bins_valid 设为真    # 重建图片同理
    end = time.time()
    print('Total time:', end - begin)